# ImportError: cannot import name 'Row' from 'sqlalchemy' (/home/legend/.local/lib/python3.10/site-packages/sqlalchemy/__init__.py)
# 解决办法是pip install --upgrade --force-reinstall flask-sqlalchemy sqlalchemy
from langchain.cache import InMemoryCache, SQLiteCache
from langchain.chains.summarize import load_summarize_chain
from langchain.globals import set_llm_cache
from langchain.llms.openai import OpenAI
from langchain.schema import Document
from langchain.text_splitter import CharacterTextSplitter

llm = OpenAI(model_name="text-davinci-002", n=2, best_of=2)


def use_memory_cache():
    set_llm_cache(InMemoryCache())
    print(llm.predict("Tell me a  joke"))


def use_sqlite_cache():
    set_llm_cache(SQLiteCache(database_path=".langchain.db"))
    print(llm.predict("Tell me a joke"))


def no_cache():
    no_cache_llm = OpenAI(model_name="text-davinci-002", cache=False)
    text_splitter = CharacterTextSplitter()
    with open("state_of_the_union.txt") as f:
        state_of_union = f.read()
    texts = text_splitter.split_text(state_of_union)
    docs = [Document(page_content=t) for t in texts[:3]]
    chain = load_summarize_chain(llm, chain_type="map_reduce", reduce_llm=no_cache_llm)
    print(chain.run(docs))


if __name__ == '__main__':
    no_cache()